const https = require('https');
const fs = require('fs');
const path = require('path');

const MODEL_URL = 'https://raw.githubusercontent.com/justadudewhohacks/face-api.js/master/weights';
const MODELS = [
  {
    name: 'ssd_mobilenetv1_model',
    files: [
      'ssd_mobilenetv1_model-weights_manifest.json',
      'ssd_mobilenetv1_model-shard1',
      'ssd_mobilenetv1_model-shard2'
    ]
  },
  {
    name: 'face_landmark_68_model',
    files: [
      'face_landmark_68_model-weights_manifest.json',
      'face_landmark_68_model-shard1'
    ]
  },
  {
    name: 'face_recognition_model',
    files: [
      'face_recognition_model-weights_manifest.json',
      'face_recognition_model-shard1',
      'face_recognition_model-shard2'
    ]
  }
];

// 确保模型目录存在
const modelDir = path.join(__dirname, 'models');
if (!fs.existsSync(modelDir)) {
  fs.mkdirSync(modelDir, { recursive: true });
}

// 下载文件函数（带重试机制）
async function downloadFile(filename, retries = 3) {
  const filePath = path.join(modelDir, filename);
  const fileUrl = `${MODEL_URL}/${filename}`;

  // 如果文件已存在且大小不为0，跳过下载
  if (fs.existsSync(filePath) && fs.statSync(filePath).size > 0) {
    console.log(`文件已存在且有效，跳过下载: ${filename}`);
    return;
  }

  for (let attempt = 1; attempt <= retries; attempt++) {
    try {
      console.log(`开始下载: ${filename} (尝试 ${attempt}/${retries})`);
      
      const response = await new Promise((resolve, reject) => {
        https.get(fileUrl, (res) => {
          if (res.statusCode === 302 || res.statusCode === 301) {
            // 处理重定向
            https.get(res.headers.location, resolve).on('error', reject);
          } else if (res.statusCode === 200) {
            resolve(res);
          } else {
            reject(new Error(`下载失败: HTTP ${res.statusCode}`));
          }
        }).on('error', reject);
      });

      const file = fs.createWriteStream(filePath);
      
      await new Promise((resolve, reject) => {
        response.pipe(file);
        
        file.on('finish', () => {
          file.close();
          // 验证文件大小
          const stats = fs.statSync(filePath);
          if (stats.size === 0) {
            fs.unlinkSync(filePath);
            reject(new Error('下载的文件大小为0'));
          } else {
            console.log(`下载完成: ${filename} (${(stats.size / 1024 / 1024).toFixed(2)}MB)`);
            resolve();
          }
        });
        
        file.on('error', (err) => {
          fs.unlink(filePath, () => {});
          reject(err);
        });
      });

      return; // 下载成功，退出函数
    } catch (error) {
      console.error(`下载 ${filename} 失败 (尝试 ${attempt}/${retries}):`, error.message);
      if (attempt === retries) {
        throw error; // 所有重试都失败了
      }
      // 等待一段时间后重试
      await new Promise(resolve => setTimeout(resolve, 2000 * attempt));
    }
  }
}

// 下载所有模型文件
async function downloadAllModels() {
  console.log('开始下载模型文件...');
  console.log('模型目录:', modelDir);
  
  let successCount = 0;
  let failCount = 0;
  
  for (const model of MODELS) {
    console.log(`\n处理模型: ${model.name}`);
    for (const file of model.files) {
      try {
        await downloadFile(file);
        successCount++;
      } catch (error) {
        console.error(`下载 ${file} 最终失败:`, error.message);
        failCount++;
      }
    }
  }
  
  console.log('\n下载完成统计:');
  console.log(`成功: ${successCount} 个文件`);
  console.log(`失败: ${failCount} 个文件`);
  
  if (failCount > 0) {
    console.log('\n警告: 部分文件下载失败，请检查网络连接后重试');
  } else {
    console.log('\n所有模型文件下载成功！');
  }
}

// 执行下载
console.log('开始下载人脸识别模型文件...');
downloadAllModels().catch(error => {
  console.error('下载过程中发生错误:', error);
  process.exit(1);
}); 